{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://csdms.colorado.edu\"><img style=\"float: center; width: 75%\" src=\"../../media/ESPIn.png\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Landlab: Grids, Landlab as a solver for advection-diffusion problems, some demos\n",
    "\n",
    "This tutorial illustrates how you can use Landlab to construct a simple two-dimensional numerical model on a regular (raster) grid, using a simple forward-time, centered-space / upwind numerical scheme. \n",
    "\n",
    "The numerical examples we have covered so far are all pretty simple. If this was all you needed to do, you wouldn't need Landlab. \n",
    "\n",
    "But what if you wanted...\n",
    "\n",
    "... to use the same diffusion model in 2D instead of 1D.\n",
    "\n",
    "... to use an irregular grid (in 1 or 2D). \n",
    "\n",
    "... wanted to combine the diffusion model with a more complex model. \n",
    "\n",
    "... have a more complex model you want to use over and over again with different boundary conditions.\n",
    "\n",
    "These are the sorts of problems that Landlab was designed to solve. \n",
    "\n",
    "In Part 1 we will use the RasterModelGrid, fields, and a numerical utility for calculating flux divergence. \n",
    "\n",
    "Landlab is often times associated with terrestrial processes. But the library provides a very rich ecosystem to tackle all kind of geophysical problems (other than the 'land'). To showcase this, we will solve the same advection and diffusion equations we solved before, but this time in 2 dimensions and using the landlab toolkit to build a model.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Landlab's Model Grids\n",
    "\n",
    "The Landlab model grids are data structures that represent the model domain (the variable `x` in our prior example). Here we will use `RasterModelGrid` which creates a grid with regularly spaced square grid elements. The RasterModelGrid knows how the elements are connected and how far apart they are.\n",
    "\n",
    "Lets start by creating a RasterModelGrid class. First we need to import it. \n",
    "\n",
    "### What are landlab grids (the gridding engine):\n",
    " - Two-dimensional simulation grid of a user-specified size and shape\n",
    " - Grids are represented as Python objects\n",
    "     - data (geomtery, topology)\n",
    "     - methods (functions to perform operations using the data)\n",
    " - Landlab mainly supports 2D grids \n",
    " \n",
    "### What you need to know about grids: \n",
    "\n",
    "<img src=\"../../media/Grids1.png\"/>\n",
    "\n",
    "**Figure** Geometry and topology of grid elements on various Landlab grids ([Hobley et al. 2017](https://esurf.copernicus.org/articles/5/21/2017/))\n",
    "\n",
    "- a set of primitive elements (see figure)\n",
    "- the entire grid can be generated from a description of the geometry of only one of these element types – typically, a user might specify the locations of the nodes, and the grid object’s remaining elements are automatically placed according to this node framework.\n",
    "- bounding elements are a set of nodes (grids are finite)\n",
    "        => A grid has nodes than cells\n",
    "\n",
    "-- Landlab provides support for regular and irregular grids. In the following examples, we will work with regular grids. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (a) Explore the RasterModelGrid\n",
    "\n",
    "Before we use a RasterModelGrid to solve our example, lets explore it a bit closer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from landlab import RasterModelGrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> RasterModelGrid refers to a `class` (OOP)\n",
    "\n",
    "We'll start with a 3-row by 4-column raster grid, with 10-meter node spacing by creating an instance of the class RastermodelGrid\n",
    "\n",
    "#### First, the nodes\n",
    "\n",
    "> How many cells has this grid? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from landlab.plot.graph import plot_graph\n",
    "grid = RasterModelGrid((3, 4), 10.0)\n",
    "plot_graph(grid, at=\"node\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that the nodes are points and they are numbered with unique IDs from lower left to upper right. \n",
    "\n",
    "Not sure how to use he RasterModelGrid class? \n",
    "Bring up the documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?RasterModelGrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that the nodes are points and they are numbered with unique IDs from lower left to upper right. \n",
    "\n",
    "Create a scalar field called z, representing elevation of the land surface, at the grid nodes:\n",
    "\n",
    "Next we'll add a data field to the grid, to represent the elevation values at grid nodes. The \"dot\" syntax below indicates that we are calling a function (or method) that belongs to the RasterModelGrid class, and will act on data contained in mg. The arguments indicate that we want the data elements attached to grid nodes (rather than links, for example), and that we want to name this data field topographic__elevation. The add_zeros method returns the newly created NumPy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = grid.add_zeros('topographic__elevation', at='node')\n",
    "z[5] = 5.0\n",
    "z[6] = 3.6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nodes 5 and 6 are the only core nodes; the rest are (so far) open boundaries.\n",
    "\n",
    "Here are the values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grid.at_node[\"topographic__elevation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a graphical look at the elevation grid we've created. To do so, we'll use the matplotlib graphics library (imported under the name plt). We also have to tell the Jupyter Notebook to display plots right here on the page. Finally, we will import Landlab's imshow_grid function to display our gridded values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from landlab.plot.imshow import imshow_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow_grid(grid, 'topographic__elevation')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are elevation values associated with all 12 of the nodes on the grid. The ones shown in black are boundary nodes; the two in the middle are core nodes. This is our scalar field.\n",
    "\n",
    "> Note: you can assign multiple fields to a node! You can make as many fields as you like and assign it to any element of the grid (nodes/links etc). Try to add a second field to the nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid.add_zeros('soil_depth',at='node')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change value in soil depth grid (class)\n",
    "# Show referncing to memory when chaning z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Second, the links\n",
    "\n",
    "We want to find the gradient between each pair of adjacent nodes, and store that value at the associated **link** between them.\n",
    "\n",
    "##### What are links?\n",
    "\n",
    "For each pair of adjacent nodes in a Landlab grid, there is a corresponding **link**. Links are directed line segments whose endpoints are the coordinates of the nodes. A link can be illustrated like this:\n",
    "\n",
    "    o---->o\n",
    "\n",
    "Here, each o represents a node, and the arrow is the link. A \"vertical\" link looks like this:\n",
    "\n",
    "    o\n",
    "    ^\n",
    "    |\n",
    "    |\n",
    "    o\n",
    "\n",
    "The tip of the arrow is known as the **link head**; the base of the arrow is known as the **link tail**. By default, links always \"point\" within the upper-right half-plane.\n",
    "\n",
    "With this definition of links in mind, we can sketch our grid like so, with the ID numbers of the nodes shown by the numbers:\n",
    "\n",
    "\n",
    "    8 ----> 9 ----> 10----> 11\n",
    "    ^       ^       ^       ^\n",
    "    |       |       |       |\n",
    "    |       |       |       |\n",
    "    4 ----> 5 ----> 6 ----> 7\n",
    "    ^       ^       ^       ^\n",
    "    |       |       |       |\n",
    "    |       |       |       |\n",
    "    0 ----> 1 ----> 2 ----> 3\n",
    "\n",
    "\n",
    "If we label each node with its elevation value, it looks like this:\n",
    "\n",
    "\n",
    "    0 ----> 0 ----> 0 ----> 0\n",
    "    ^       ^       ^       ^\n",
    "    |       |       |       |\n",
    "    |       |       |       |\n",
    "    0 ---->5.0---->3.6----> 0\n",
    "    ^       ^       ^       ^\n",
    "    |       |       |       |\n",
    "    |       |       |       |\n",
    "    0 ----> 0 ----> 0 ----> 0\n",
    "    \n",
    "\n",
    "Let's plot the layout of the links of our ealier created grid object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_graph(grid, at=\"link\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which are lines that connect the nodes and each have a unique ID number. \n",
    "\n",
    "We can also plot the cells which are polygons centered around the nodes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_graph(grid, at=\"cell\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Landlab is a \"dual\" graph because it also keeps track of a second set of points, lines, and polygons (\"corners\", \"faces\", and \"patches\"). We will not focus on them further.\n",
    "\n",
    "An overview ([Hobley et al. 2017](https://esurf.copernicus.org/articles/5/21/2017/)): \n",
    "<img src=\"../../media/Grids3.png\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Exercises for section 2a*\n",
    "\n",
    "(2a.1) Create an instance of a `RasterModelGrid` with 5 rows and 7 columns, with a spacing between nodes of 10 units. Plot the node layout, and identify the ID number of the center-most node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2a.1 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2a.2) Find the ID of the cell that contains this node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2a.2 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2a.3) Find the ID of the horizontal link that connects to the last node on the right in the middle column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2a.3 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "rmg.links_at_node[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gradients\n",
    "To calculate the flux of particles, we need to know the gradient gradient of a node field, with one gradient value for each link. Calculate the gradient for the grid defined in part  1 by using the function calc_grad_at_link:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx = grid.calc_grad_at_link(z)\n",
    "dzdx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a crude graphical representation of gradient array:\n",
    "\n",
    "\n",
    "    o ---0--> o ---0--> o ---0--> o\n",
    "    ^         ^         ^         ^\n",
    "    0       -0.5      -0.36       0\n",
    "    |         |         |         |\n",
    "    o  +0.5 > o -0.14 > o -0.36 > o\n",
    "    ^         ^         ^         ^\n",
    "    0       +0.5      +0.36       0\n",
    "    |         |         |         |\n",
    "    o ---0--> o ---0--> 0 ---0--> 0\n",
    "\n",
    "Links are listed in order by the $(x, y)$ coordinates of their midpoints. The ID numbering scheme for our links looks like this:\n",
    "\n",
    "\n",
    "    o --14--> o --15--> o --16--> o\n",
    "    ^         ^         ^         ^\n",
    "    10       11        12        13\n",
    "    |         |         |         |\n",
    "    o ---7--> o ---8--> o ---9--> o\n",
    "    ^         ^         ^         ^\n",
    "    3         4         5         6\n",
    "    |         |         |         |\n",
    "    o ---0--> o ---1--> 0 ---2--> 0\n",
    "\n",
    "Let's explore how the geometry and the values in the ID array of gradients correspond. Here are the gradients first three are the horizontal links along the bottom edge of the grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next come four vertical links that connect the bottom to the middle rows of nodes. Two of these values are positive, indicating an uphill gradient in the direction of the links:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[3:7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the middle row of horizontal links:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[7:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next row of vertical links. The middle two of these are negative, indicating a downhill slope in the direction of the links:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[10:14]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the top row of horizontal links:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[14:17]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An alternative way to inspect link-based values in a raster grid is to use the horizontal_links and vertical_links grid attributes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[grid.horizontal_links]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dzdx[grid.horizontal_links]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Exercises for section 2c*\n",
    "\n",
    "(2c.1) Make a 3x3 `RasterModelGrid` called `tinygrid`, with a cell spacing of 2 m. Use the `plot_graph` function to display the nodes and their ID numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.1 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2c.2) Give your `tinygrid` a node field called `height` and set the height of the center-most node to 0.5. Use `imshow_grid` to display the height field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.2 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2c.3) The grid should have 12 links (extra credit: verify this with `plot_graph`). When you compute gradients, which of these links will have non-zero gradients? What will the absolute value(s) of these gradients be? Which (if any) will have positive gradients and which negative? To codify your answers, make a 12-element numpy array that contains your predicted gradient value for each link."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.3 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2c.4) Test your prediction by running the `calc_grad_at_link` function on your tiny grid. Print the resulting array and compare it with your predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.4 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2c.5) Suppose the flux of soil per unit cell width is defined as -0.01 times the height gradient. What would the flux be at the those links that have non-zero gradients? Test your prediction by creating and printing a new array whose values are equal to -0.01 times the link-gradient values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.5 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2c.6) Consider the net soil accumulation or loss rate around the center-most node in your tiny grid (which is the only one that has a cell). The *divergence* of soil flux can be represented numerically as the sum of the total volumetric soil flux across each of the cell's four faces. What is the flux across each face? (Hint: multiply by face width) What do they add up to? Test your prediction by running the grid function `calc_flux_div_at_node` (hint: pass your unit flux array as the argument). What are the units of the divergence values returned by the `calc_flux_div_at_node` function?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2c.6 here)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Exercises for section 2b*\n",
    "\n",
    "#### Preparation: \n",
    "\n",
    "load the map of Europe, plot the location of the eyjafjallajökull and Brussels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EU_coast = np.loadtxt('../../data/EuropePoints_xp25_ym30.csv', delimiter=',')\n",
    "vol_y = 35\n",
    "vol_x = 7\n",
    "Br_y = 21\n",
    "Br_x = 30\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(EU_coast[:,0],EU_coast[:,1],s=0.5,c='r')\n",
    "plt.scatter(vol_x,vol_y,s=5,c='k')\n",
    "plt.scatter(Br_x,Br_y,s=5,c='g')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b .1) Create an instance of a `RasterModelGrid` called `EU`, with 100 rows and 140 columns, with a spacing between nodes of 0.5 degrees (1 degree = ca. 111 km). Calculate the index of the eyjafjallajökull and Brussels within the grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2b.1 here)\n",
    "rows = 100\n",
    "cols = 140\n",
    "dx = 0.5 # in degrees\n",
    "# xy spacing in degrees (ca. 111 km)\n",
    "EU = RasterModelGrid((rows, cols), xy_spacing=dx)\n",
    "\n",
    "vol_loc = int(vol_y/dx*cols+vol_x/dx)\n",
    "Br_loc = int(Br_y/dx*cols+Br_x/dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b.2) Query the grid variables `number_of_nodes` and `number_of_core_nodes` to find out how many nodes are in your grid, and how many of them are core nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2b.2 here)\n",
    "print(EU.number_of_nodes)\n",
    "print(EU.number_of_core_nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b.3) Add a new field to your grid, called `Aerosol` and attached to nodes. Have the initial values be all zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2b.3 here)\n",
    "a = EU.add_zeros('Aerosol', at='node')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b.4) Change the Aerosol concentration at the eyjafjallajökull to an initial concentration of 10 ppm (store this value in `C_ini`) . Use the `imshow_grid` function to display a shaded image of the Aerosol field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (enter your solution to 2b.4 here)\n",
    "C_ini = 10.0\n",
    "a[vol_loc] = C_ini\n",
    "imshow_grid(EU, 'Aerosol')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b.5) Now lets do some more advanced plotting. Use the `imshow_grid` function to to plot the Aerosol concentration, but this time plot all cells with Aerosol concentrations smaller than 1 ppm as white transparent pixels. To do this, make a copy of the Aerosol field and mask it [see ](https://numpy.org/doc/stable/reference/maskedarray.generic.html). Hint: try using `np.ma.masked_where(..., ...)`. Use the inferno_r colormap (see cmap argument of `imshow_grid`). On top of the concentration field, plot the map of Europe using the statements from the preparation cell above and the EU_coast variable. \n",
    "\n",
    "Finally, make a function out of this plotting code snippet so that you can easily reuse it later "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_aerosol(a):\n",
    "    a_plt = np.zeros_like(a)\n",
    "    a_plt[:] = a\n",
    "    a_plt = np.ma.masked_where(a_plt <1, a_plt)\n",
    "    imshow_grid(EU,a_plt,cmap ='inferno_r',color_for_background=None, color_for_closed=None)        \n",
    "    plt.scatter(EU_coast[:,0],EU_coast[:,1],s=0.5,c='r')\n",
    "    plt.scatter(vol_x,vol_y,s=5,c='k')\n",
    "    plt.scatter(Br_x,Br_y,s=5,c='g')        \n",
    "    plt.show()\n",
    "    \n",
    "plot_aerosol(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some constants: \n",
    "- $C_n$ (source term): amount of dust particles produced per day (20 ppm/day)\n",
    "- $D$ (diffusion constant) : 1.5  deg$^2$  day$^-1$\n",
    "- $dt$ (in days, see stability rule before)\n",
    "- get $dx$ and $dy$ from landlab grid\n",
    "- totT: the total amount of iterations, see to 1000 for now\n",
    "- create an empty np array to store the aerosol concentration above Brussels during the model run (with size totT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_n = 200\n",
    "D  = 1.5\n",
    "dx = EU.dx\n",
    "dy = EU.dy\n",
    "dt= 0.245 * EU.dx * EU.dx / D\n",
    "totT = 1000\n",
    "Br_out = np.zeros(totT)\n",
    "time = np.arange(start=0,stop = totT, step = 1)*dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "- Create a new field of zeros called aerosol_flux and attached to links. Using the 'number_of_links' grid property, verify that your new field array has the correct number of items. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa = EU.add_zeros('aerosol_flux', at='link')\n",
    "print(EU.number_of_links)\n",
    "print(len(qa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Boundary conditions: for this example, we'll assume that all sides are  are open. (The order of the function arguments is east, north, west, south). This is the default configuration. LAter on we will check what happens if boundary conditions are open. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EU.set_closed_boundaries_at_grid_edges(True, True, True, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*A note on boundaries:* with a Landlab raster grid, all the perimeter nodes are boundary nodes. In this example, there are 98 + 140 + 98 + 140 = 476 boundary nodes. All the remaining nodes are known as **core** nodes. In this example, there are 100*140-476 = 13524 core nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(EU.boundary_nodes)\n",
    "# len(EU.core_nodes)\n",
    "100*140-476\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(2b.7) Now solve the 2D diffusion equation. Loop through totT iterations, representing $totT*dt$ days. On each pass through the loop, we do the following:\n",
    "\n",
    "1. Calculate, and store in the array `g`, the **aerosol gradient** between each neighboring pair of nodes. The gradient value is a positive number when the gradient is \"uphill\" in the direction of the link, and negative when the gradient is \"downhill\" in the direction of the link. On a raster grid, link directions are always in the direction of increasing 𝑥 (\"horizontal\" links) or increasing 𝑦 (\"vertical\" links).\n",
    "\n",
    "2. Calculate, and store in the array `qa`, the **aerosol flux** between each adjacent pair of nodes by multiplying their gradient by the transport coefficient. We will only do this for the active links (those not connected to a closed boundary, and not connecting two boundary nodes of any type); others will remain as zero.\n",
    "\n",
    "3. Calculate the resulting net flux at each node, (positive=net outflux, negative=net influx). The negative of this array is the rate of change of elevation at each (core) node, so store it in a node array called `dadt'.\n",
    "\n",
    "4. Update the aerosol concentrations for the new time step.\n",
    "\n",
    "5. Take care of the source term: calculate the amount of ash being produced during one timestep and add this to the aerosol field at the volcano location\n",
    "\n",
    "5. Store the concentration above Brussels in the `Br_out` array\n",
    "\n",
    "6. Plot the resulting concentration every 100 iterations using the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in range(totT):  \n",
    "#     g = \n",
    "#     qa[EU.active_links] = \n",
    "#     dadt = \n",
    "#     a[EU.core_nodes] += dadt[EU.core_nodes] * dt   \n",
    "#     a[vol_loc] += C_n*dt\n",
    "#     Br_out[i] = a[Br_loc]       \n",
    "#     if i%100==0:     \n",
    "#         print('Time is : ' + str(i*dt) + 'days')\n",
    "#         plot_aerosol(a)      \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the concentration through time at Brussels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(time,Br_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* What happens if you set all of the boundaries as closed boundaries? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advection diffusion: \n",
    "We have solved the diffusion part of the volcanic ash problem imposed by the eyjafjallajökull volcano. However, most of the ash transport was induced by wind fields advection the ash towards Europe (yup, west winds). Can you come up with a landlab implementation of the advection equation? Copy paste your solution from exercise () and add some lines to solve the 2D advection equation:\n",
    "\n",
    "$$\\frac{\\partial C}{\\partial t} +v \\frac{\\partial C}{\\partial x} +u \\frac{\\partial C}{\\partial y}=0\n",
    "\\label{eq:1}\\tag{1}$$\n",
    "where $C$ is the aerosol concentration and $v$ is a constant windspeed at which aerosol concentrations are advectected in the x direction and $u$ is a constant windspeed at which aerosol concentrations are advectected in the y direction.\n",
    "\n",
    "You will need some additional variables: \n",
    "- $v$ (horizontal wind speed, positive in x direction): 1.5 deg/day\n",
    "- $u$ (vertical wind speed, positive in y direction): -0.75 deg/day\n",
    "\n",
    "- Hint 1: reset the Aerosol field (a)\n",
    "- Hint 2: update the time criterion and combine the diffusion and the advection criteria \n",
    "- Hint 3: currently, the landlab grid library has no predefined functions to calculate advection (no equivalent for calc_flux_div_at_node which is used to solve the heat eq.). Hence, you will have to discretise the advection equation yourself using an upwind first order FDM, but using the matlab node/link structure. To calculate the links at every node, use `links_at_node`\n",
    "- Hint 4: given that there is no predefined method to calculate advection, you will also have to take care of the boundary conditions. Assume open boundary nodes by setting the aerosol concentration at all boundary nodes to 0 ppm at the end of every iteration. Use `grid.boundary_nodes`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v = 1.5\n",
    "# u = -.5\n",
    "# C_n = 200\n",
    "# D  = 1.5\n",
    "# dx = EU.dx\n",
    "# dy = EU.dy\n",
    "# totT = 1000\n",
    "# Br_out = np.zeros(totT)\n",
    "# time = np.arange(start=0,stop = totT, step = 1)*dt\n",
    "\n",
    "# a[:] = np.zeros_like(a)\n",
    "# g = np.zeros(EU.number_of_links)\n",
    "# dadt = np.zeros(EU.number_of_nodes)\n",
    "\n",
    "\n",
    "# dt_a = 0.95*np.minimum(dx/abs(v),dy/abs(u))\n",
    "# dt_d = 0.245 * dx * dx / D\n",
    "# dt = min(dt_a,dt_d)\n",
    "\n",
    "# rt = EU.links_at_node[0:len(a)][:,0]\n",
    "# up = EU.links_at_node[0:len(a)][:,1]\n",
    "# lt = EU.links_at_node[0:len(a)][:,2]\n",
    "# dw = EU.links_at_node[0:len(a)][:,3] \n",
    "\n",
    "# for i in range(totT):  \n",
    "#     EU.calc_grad_at_link(a,out=g)\n",
    "#     qa[EU.active_links] = ...\n",
    "#     EU.calc_flux_div_at_node(qa,out=dadt)\n",
    "#     a[EU.core_nodes] -= ...\n",
    "    \n",
    "#     # Advection \n",
    "#     # horizontal (x) direction\n",
    "#     EU.calc_grad_at_link(a,out=g)    \n",
    "#     if v<0:\n",
    "#         a -= v*g[rt]*dt\n",
    "#     else:             \n",
    "#         a -= ...\n",
    "    \n",
    "#     # vertical (y) direction\n",
    "#     EU.calc_grad_at_link(a,out=g)\n",
    "#     if u<0:\n",
    "#         a -= ...\n",
    "#     else:\n",
    "#         a -= u*g[dw]*dt\n",
    "    \n",
    "#     # Or shorter: \n",
    "# #     EU.calc_grad_at_link(a,out=g)   \n",
    "# #     a -= (v*g[rt]*(v<0) + v*g[lt]*(v>0))*dt\n",
    "# #     EU.calc_grad_at_link(a,out=g)\n",
    "# #     a -= (u*g[up]*(u<0) + u*g[dw]*(u>0))*dt\n",
    "    \n",
    "#     # Source term   \n",
    "#     a[vol_loc] += C_n*dt\n",
    "    \n",
    "#     # Keep track of concentration in Brussels\n",
    "#     Br_out[i] = a[Br_loc]\n",
    "    \n",
    "#     #BC    \n",
    "#     a[mg.boundary_nodes] = 0 \n",
    "#     if i%100==0:     \n",
    "#         print('Time is : ' + str(i*dt) + 'days')\n",
    "#         plot_aerosol(a) \n",
    "        \n",
    "# plt.figure()\n",
    "# plt.plot(time,Br_out) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* what happens if you change the vertical (u) velocity to  +0.25? \n",
    "\n",
    "Congratulations on making it to the end of this tutorial!\n",
    "\n",
    "### Click here for more <a href=\"https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html\">Landlab tutorials</a>"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
